import torch
import torch.nn as nn
# from timm.models import convmixer


class Residual(nn.Module):
    def __init__(self, fn):
        super().__init__()
        self.fn = fn

    def forward(self, x):
        return self.fn(x) + x


class ConvMixer(nn.Module):
    def __init__(
            self,
            dim,
            depth,
            kernel_size=9,
            patch_size=7,
            in_chans=3,
            num_classes=1000,
            global_pool='avg',
            drop_rate=0.,
            act_layer=nn.GELU,
            **kwargs,
    ):
        super().__init__()
        self.num_classes = num_classes
        self.num_features = dim
        self.grad_checkpointing = False

        self.stem = nn.Sequential(
            nn.Conv2d(in_chans, dim, kernel_size=patch_size, stride=patch_size),
            act_layer(),
            nn.BatchNorm2d(dim)
        )
        self.blocks = nn.Sequential(
            *[nn.Sequential(
                Residual(nn.Sequential(
                    nn.Conv2d(dim, dim, kernel_size, groups=dim, padding="same"),
                    act_layer(),
                    nn.BatchNorm2d(dim)
                )),
                nn.Conv2d(dim, dim, kernel_size=1),
                act_layer(),
                nn.BatchNorm2d(dim)
            ) for i in range(depth)]
        )
        # self.pooling = SelectAdaptivePool2d(pool_type=global_pool, flatten=True)
        self.pooling = nn.AdaptiveAvgPool2d((1, 1))
        self.flatten = nn.Flatten()
        self.head_drop = nn.Dropout(drop_rate)
        self.head = nn.Linear(dim, num_classes) if num_classes > 0 else nn.Identity()


    def forward_features(self, x):
        x = self.stem(x)
        x = self.blocks(x)
        return x

    def forward_head(self, x, pre_logits: bool = False):
        x = self.pooling(x)
        x = self.flatten(x)
        x = self.head_drop(x)
        return x if pre_logits else self.head(x)

    def forward(self, x):
        x = self.forward_features(x)
        x = self.forward_head(x)
        return x


def convmixer_1536_20(**kwargs) -> ConvMixer:
    model = ConvMixer(dim=1536, depth=20, kernel_size=9, patch_size=7, **kwargs)
    return model

def convmixer_768_32(**kwargs) -> ConvMixer:
    model = ConvMixer(dim=768, depth=32, kernel_size=7, patch_size=7, act_layer=nn.ReLU, **kwargs)
    return model


def convmixer_1024_20_ks9_p14(**kwargs) -> ConvMixer:
    model = ConvMixer(dim=1024, depth=20, kernel_size=9, patch_size=14, **kwargs)
    return model


if __name__ == '__main__':
    model = convmixer_1024_20_ks9_p14()
    model.eval()
    out = model(torch.rand(1, 3, 224, 224))
    print(out)
